﻿using System;
using System.IO;
using System.Security.Cryptography;
using System.Text;

namespace CS795.TravelAgency.Encryption
{
    public static class AESEncryption
    {
        public static string Encrypt(string plainText, string passphrase, string salt = "12345678", int keySize = 256)
        {
            if (string.IsNullOrEmpty(plainText))
                throw new ArgumentNullException("plainText");

            if (string.IsNullOrEmpty(passphrase))
                throw new ArgumentNullException("passphrase");

            string result = "";

            RijndaelManaged aes = new RijndaelManaged();
            Rfc2898DeriveBytes key = new Rfc2898DeriveBytes(passphrase, Encoding.UTF8.GetBytes(salt));
            aes.KeySize = keySize;
            aes.Key = key.GetBytes(aes.KeySize / 8);
            using (MemoryStream mem = new MemoryStream())
            {
                mem.Write(BitConverter.GetBytes(aes.IV.Length), 0, sizeof(int));
                mem.Write(aes.IV, 0, aes.IV.Length);
                using (ICryptoTransform encryptor = aes.CreateEncryptor(aes.Key, aes.IV))
                {
                    using (CryptoStream cs = new CryptoStream(mem, encryptor, CryptoStreamMode.Write))
                    {
                        using (StreamWriter writer = new StreamWriter(cs))
                        {
                            writer.Write(plainText);
                        }
                    }
                    result = Convert.ToBase64String(mem.ToArray());
                }
            }
            return result;
        }

        public static string Decrypt(string cipherText, string passphrase, string salt = "12345678", int keySize = 256)
        {
            if (string.IsNullOrEmpty(cipherText))
                throw new ArgumentNullException("cipherText");

            if (string.IsNullOrEmpty(passphrase))
                throw new ArgumentNullException("passphrase");

            string result = "";

            RijndaelManaged aes = new RijndaelManaged();
            Rfc2898DeriveBytes key = new Rfc2898DeriveBytes(passphrase, Encoding.UTF8.GetBytes(salt));
            aes.KeySize = keySize;
            aes.Key = key.GetBytes(aes.KeySize / 8);
            using (MemoryStream mem = new MemoryStream(Convert.FromBase64String(cipherText)))
            {
                byte[] rawLength = new byte[sizeof(int)];
                if (mem.Read(rawLength, 0, rawLength.Length) != rawLength.Length)
                {
                    throw new SystemException("Stream did not contain properly formatted byte array");
                }

                byte[] buffer = new byte[BitConverter.ToInt32(rawLength, 0)];
                if (mem.Read(buffer, 0, buffer.Length) != buffer.Length)
                {
                    throw new SystemException("Did not read byte array properly");
                }

                aes.IV = buffer;
                using (ICryptoTransform decryptor = aes.CreateDecryptor(aes.Key, aes.IV))
                {
                    using (CryptoStream cs = new CryptoStream(mem, decryptor, CryptoStreamMode.Read))
                    {
                        using (StreamReader reader = new StreamReader(cs))
                        {
                            result = reader.ReadToEnd();
                        }
                    }
                }
            }
            return result;

        }
    }
}